package class07;

import java.util.*;

/**
 * 加强堆
 */
public class MyHeapGreater<T> {

    // 将数据封装成引用类型，确保每一个Key都是唯一的
    private static class DataType<T> {
        T data;

        private DataType(T data) {
            this.data = data;
        }
    }

    private final ArrayList<DataType<T>> heap;  // 堆
    private final Map<DataType<T>, Integer> indexHeap;  // 反向索引：每个堆所在位置下标
    private final Comparator<? super DataType<T>> comp;  // 比较器规则

    public MyHeapGreater(Comparator<? super DataType<T>> comp) {
        this.heap = new ArrayList<>();
        this.indexHeap = new HashMap<>();
        this.comp = comp;
    }

    // 堆是否为空
    public boolean isEmpty() {
        return heap.size() == 0;
    }

    // 堆大小
    public int size() {
        return heap.size();
    }

    // 堆是否包含某个数据
    public boolean contains(DataType<T> data) {
        return indexHeap.containsKey(data);
    }

    // 查看堆顶数据
    public DataType<T> peek() {
        return heap.get(0);
    }

    // 向堆中添加一个数据
    public void push(DataType<T> data) {
        heap.add(data);
        indexHeap.put(data, size());
        heapInsert(0);
    }

    // 抛出堆顶数据
    public DataType<T> poll() {
        swap(0, size() - 1);
        DataType<T> result = heap.remove(size() - 1);
        indexHeap.remove(result);
        heapify(0);
        return result;
    }

    // 删除堆中某个数据，根据已有反向索引，可以 O(log n) 高效删除
    public void remove(DataType<T> data) {
        DataType<T> tempData = heap.get(size() - 1);
        int tempIndex = indexHeap.get(data);
        indexHeap.remove(data);
        heap.remove(size() - 1);
        if (data != tempData) {
            heap.set(tempIndex, tempData);
            indexHeap.put(tempData, tempIndex);
            resign(tempData);
        }
    }

    // 堆中某个数据被修改后，为确保还是堆结构，需要调用该函数，保持堆结构
    public void resign(DataType<T> data) {
        // 修改的数据要么向上，要么向下，这两个函数只会去跑一个
        heapInsert(indexHeap.get(data));
        heapify(indexHeap.get(data));
    }

    // 获取堆中所有元素
    public ArrayList<DataType<T>> getAllElements() {
        return new ArrayList<>(heap);
    }

    @SuppressWarnings("all")  // 取消提示
    private void heapify(int index) {
        int left = index * 2 + 1;
        while (left < size()) {
            int best = left + 1 < size() && comp.compare(heap.get(left + 1), heap.get(left)) < 0 ? left + 1 : left;
            best = comp.compare(heap.get(index), heap.get(best)) < 0 ? index : best;
            if (best == index) {
                break;
            }
            swap(index, best);
            index = best;
            left = index * 2 + 1;
        }
    }

    // 按堆的规则进行插入一个数据，调整堆结构
    @SuppressWarnings("all")
    private void heapInsert(int index) {
        while (comp.compare(heap.get(index), heap.get((index - 1) / 2)) < 0) {
            swap(index, (index - 1) / 2);
            index = (index - 1) / 2;
        }
    }

    private void swap(int index, int root) {
        DataType<T> t1 = heap.get(index);
        DataType<T> t2 = heap.get(root);
        indexHeap.put(t2, index);
        indexHeap.put(t1, root);
        heap.set(index, t2);
        heap.set(root, t1);
    }

}
